import numpy as np
import torch
import os
from PIL import Image
from codes.misc_fucntion import convert_to_grayscale
from codes.misc_fucntion import save_gradient_images
from codes.misc_fucntion import preprocess_image, save_class_activation_images

from codes.gradcam import GradCam
from codes.guided_backprop import GuidedBackprop


def load_model():
    model = torch.load('/home/imed/Research/TortuosityGrading/checkpoint/ResNet18(Coarse2Fine)_all_classes6666.pth')
    return model


def guided_grad_cam(grad_cam_mask, guided_backprop_mask):
    """
    Guided grad cam is just pointwise multiplication of cam mask and guided backpropagation mask
    """
    cam_gb = np.multiply(grad_cam_mask, guided_backprop_mask)
    return cam_gb


if __name__ == '__main__':
    """
    """
    original_image = "/home/imed/Research/TortuosityGrading/dataset/TortuosityLevel(1-23-4)/test/img/李媚2(47)_level3.jpg"
    seg_image = "/home/imed/Research/TortuosityGrading/dataset/TortuosityLevel(1-23-4)/test/seg/李媚2(47)_level3.png"
    img = Image.open(original_image)
    image = preprocess_image(img, resize=224)
    seg = Image.open(seg_image)
    seg = preprocess_image(seg, resize=224)

    preprocessed_img = torch.cat((image, seg), dim=1)

    target_class = None
    file_name_to_save = "/home/imed/Research/CnnVisualization/outputs/" + os.path.basename(original_image)[:-4]
    pretrained_model = load_model()
    pretrained_model = pretrained_model.module.net

    # Grad cam
    grad_cam = GradCam(pretrained_model, target_layer="layer4")
    # generate cam mask
    cam = grad_cam.generate_cam(preprocessed_img, target_class)
    save_class_activation_images(img.resize((224, 224)), cam, file_name_to_save + "_grad_cam")

    # Guided back propagation
    GBP = GuidedBackprop(pretrained_model)
    guided_grads = GBP.generate_gradients(preprocessed_img, target_class)

    # guided grad cam
    cam_gb = guided_grad_cam(cam, guided_grads)
    save_gradient_images(cam_gb, file_name_to_save + "_guided_grad_cam")
    grayscale_cam_gb = convert_to_grayscale(cam_gb)
    save_gradient_images(grayscale_cam_gb, file_name_to_save + "_guided_grad_cam_grayscale")
